"""
model name : 目标跟踪
file       : 2_YOLO数据标签目标跟踪可视化.py
information:
    author : OuYang
    time   : 2025/7/19
"""
import os

import cv2

image_dir = "E:\data\\target tracking\YOLO_OTB\Car24\\train\images"
label_dir = "E:\data\\target tracking\YOLO_OTB\Car24\\train\labels"


def show_one(image_path, label_path):
    image_src = cv2.imread(image_path)

    # 读取标签数据
    labels = []
    with open(label_path, "r") as f:

        for line in f.readlines():
            label = line.strip().split()

            x = float(label[1])
            y = float(label[2])
            w = float(label[3])
            h = float(label[4])
            labels.append([x, y, w, h])

    # 获取图片大小
    width, height = image_src.shape[1], image_src.shape[0]

    # 绘制矩形框
    for label in labels:
        x, y, w, h = label

        x1 = int((x - w / 2) * width)
        y1 = int((y - h / 2) * height)
        x2 = int((x + w / 2) * width)
        y2 = int((y + h / 2) * height)

        cv2.rectangle(image_src, (x1, y1), (x2, y2), (0, 255, 0))

    return image_src


def show_all(image_dir, label_dir):
    for i, image_name in enumerate(os.listdir(image_dir)):
        image_path = os.path.join(image_dir, image_name)
        label_path = os.path.join(label_dir, image_name.split(".")[0] + ".txt")

        # 绘制一帧
        image = show_one(image_path, label_path)

        # 绘制帧数
        image = cv2.putText(image, f'{i}', (image.shape[1] - 100, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        cv2.imshow("image", image)
        if cv2.waitKey(25) & 0xFF == ord('q'):
            break


if __name__ == "__main__":
    show_all(image_dir, label_dir)
